"""
配置加载器模块

该模块提供了统一的配置文件加载和管理功能，支持多模态推理系统的各种配置需求。
主要功能包括：
- 加载YAML格式的配置文件
- 提供类型安全的配置访问接口
- 支持配置验证和默认值设置
- 兼容原有的命令行参数系统
"""

import yaml
import os
from typing import Dict, Any, Optional
from dataclasses import dataclass, field
import logging

logger = logging.getLogger(__name__)


@dataclass
class ModelConfig:
    """模型配置类"""
    type: str = "vllm"
    vllm_config: Dict[str, Any] = field(default_factory=dict)
    openai_config: Dict[str, Any] = field(default_factory=dict)
    generation_config: Dict[str, Any] = field(default_factory=dict)


@dataclass
class DataConfig:
    """数据配置类"""
    dataset_path: str = ""
    image_root: str = ""
    max_samples: int = -1
    field_mapping: Dict[str, str] = field(default_factory=dict)


@dataclass
class MCTSConfig:
    """MCTS配置类"""
    # 基础MCTS参数
    num_rollouts: int = 8
    max_depth_allowed: int = 5
    
    # UCT算法参数
    exploration_weight: float = 1.4
    weight_scheduler: str = "const"
    discount_factor: float = 1.0
    
    # 推理行为参数 (A1-A5)
    num_a1_steps: int = 3
    num_subquestions: int = 3
    num_votes: int = 10
    num_last_votes: int = 32
    
    # 行为开关
    disable_a1: bool = False
    disable_a5: bool = False
    
    # 高级功能
    enable_potential_score: bool = False
    modify_prompts_for_rephrasing: bool = False


@dataclass
class ExperimentConfig:
    """实验配置类"""
    # 基础实验参数
    seed: int = 42
    note: str = "default"
    
    # 数据范围
    start_idx: int = 0
    end_idx: int = -1
    
    # 调试和可视化
    verbose: bool = False
    save_tree: bool = False
    save_intermediate: bool = True
    
    # 系统配置
    max_num_worker: int = 4
    test_batch_size: int = 1


@dataclass
class OutputConfig:
    """输出配置类"""
    # 输出路径
    results_path: str = "outputs/results.jsonl"
    outputs_dir: str = "./run_outputs"
    
    # 日志配置
    log_level: str = "INFO"
    
    # 结果保存配置
    save_answer_sheets: bool = True
    save_solutions: bool = True
    save_statistics: bool = True


@dataclass
class MultimodalConfig:
    """多模态配置类"""
    # 图像处理配置
    image_format: str = "base64"
    max_image_size: int = 2048
    
    # 推理配置
    always_include_image: bool = True
    image_placeholder: str = "<image1>"
    
    # 调试信息
    show_image_info: bool = False
    log_image_usage: bool = False


@dataclass
class Config:
    """主配置类"""
    model: ModelConfig = field(default_factory=ModelConfig)
    data: DataConfig = field(default_factory=DataConfig)
    mcts: MCTSConfig = field(default_factory=MCTSConfig)
    experiment: ExperimentConfig = field(default_factory=ExperimentConfig)
    output: OutputConfig = field(default_factory=OutputConfig)
    multimodal: MultimodalConfig = field(default_factory=MultimodalConfig)


def load_config(config_path: str) -> Config:
    """
    加载配置文件
    
    Args:
        config_path (str): 配置文件路径
        
    Returns:
        Config: 配置对象
        
    Raises:
        FileNotFoundError: 配置文件不存在
        yaml.YAMLError: YAML格式错误
        ValueError: 配置参数无效
    """
    if not os.path.exists(config_path):
        raise FileNotFoundError(f"配置文件不存在: {config_path}")
    
    logger.info(f"加载配置文件: {config_path}")
    
    with open(config_path, 'r', encoding='utf-8') as f:
        config_data = yaml.safe_load(f)
    
    # 创建配置对象
    config = Config()
    
    # 加载模型配置
    if 'model' in config_data:
        model_data = config_data['model']
        config.model = ModelConfig(
            type=model_data.get('type', 'vllm'),
            vllm_config=model_data.get('vllm', {}),
            openai_config=model_data.get('openai_api', {}),
            generation_config=model_data.get('generation', {})
        )
    
    # 加载数据配置
    if 'data' in config_data:
        data_data = config_data['data']
        config.data = DataConfig(
            dataset_path=data_data.get('dataset_path', ''),
            image_root=data_data.get('image_root', ''),
            max_samples=data_data.get('max_samples', -1),
            field_mapping=data_data.get('field_mapping', {})
        )
    
    # 加载MCTS配置
    if 'mcts' in config_data:
        mcts_data = config_data['mcts']
        config.mcts = MCTSConfig(
            num_rollouts=mcts_data.get('num_rollouts', 8),
            max_depth_allowed=mcts_data.get('max_depth_allowed', 5),
            exploration_weight=mcts_data.get('exploration_weight', 1.4),
            weight_scheduler=mcts_data.get('weight_scheduler', 'const'),
            discount_factor=mcts_data.get('discount_factor', 1.0),
            num_a1_steps=mcts_data.get('num_a1_steps', 3),
            num_subquestions=mcts_data.get('num_subquestions', 3),
            num_votes=mcts_data.get('num_votes', 10),
            num_last_votes=mcts_data.get('num_last_votes', 32),
            disable_a1=mcts_data.get('disable_a1', False),
            disable_a5=mcts_data.get('disable_a5', False),
            enable_potential_score=mcts_data.get('enable_potential_score', False),
            modify_prompts_for_rephrasing=mcts_data.get('modify_prompts_for_rephrasing', False)
        )
    
    # 加载实验配置
    if 'experiment' in config_data:
        exp_data = config_data['experiment']
        config.experiment = ExperimentConfig(
            seed=exp_data.get('seed', 42),
            note=exp_data.get('note', 'default'),
            start_idx=exp_data.get('start_idx', 0),
            end_idx=exp_data.get('end_idx', -1),
            verbose=exp_data.get('verbose', False),
            save_tree=exp_data.get('save_tree', False),
            save_intermediate=exp_data.get('save_intermediate', True),
            max_num_worker=exp_data.get('max_num_worker', 4),
            test_batch_size=exp_data.get('test_batch_size', 1)
        )
    
    # 加载输出配置
    if 'output' in config_data:
        output_data = config_data['output']
        config.output = OutputConfig(
            results_path=output_data.get('results_path', 'outputs/results.jsonl'),
            outputs_dir=output_data.get('outputs_dir', './run_outputs'),
            log_level=output_data.get('log_level', 'INFO'),
            save_answer_sheets=output_data.get('save_answer_sheets', True),
            save_solutions=output_data.get('save_solutions', True),
            save_statistics=output_data.get('save_statistics', True)
        )
    
    # 加载多模态配置
    if 'multimodal' in config_data:
        mm_data = config_data['multimodal']
        config.multimodal = MultimodalConfig(
            image_format=mm_data.get('image_format', 'base64'),
            max_image_size=mm_data.get('max_image_size', 2048),
            always_include_image=mm_data.get('always_include_image', True),
            image_placeholder=mm_data.get('image_placeholder', '<image1>'),
            show_image_info=mm_data.get('show_image_info', False),
            log_image_usage=mm_data.get('log_image_usage', False)
        )
    
    # 验证配置
    _validate_config(config)
    
    logger.info("配置加载完成")
    return config


def _validate_config(config: Config):
    """
    验证配置参数的有效性
    
    Args:
        config (Config): 配置对象
        
    Raises:
        ValueError: 配置参数无效
    """
    # 验证模型类型
    if config.model.type not in ['vllm', 'openai_api']:
        raise ValueError(f"不支持的模型类型: {config.model.type}")
    
    # 验证MCTS参数
    if config.mcts.num_rollouts <= 0:
        raise ValueError("num_rollouts必须大于0")
    
    if config.mcts.max_depth_allowed <= 0:
        raise ValueError("max_depth_allowed必须大于0")
    
    if config.mcts.exploration_weight < 0:
        raise ValueError("exploration_weight必须非负")
    
    if config.mcts.weight_scheduler not in ['const', 'exp', 'lin']:
        raise ValueError(f"不支持的权重调度策略: {config.mcts.weight_scheduler}")
    
    # 验证数据配置
    if not config.data.dataset_path:
        raise ValueError("dataset_path不能为空")
    
    # 验证输出配置
    if not config.output.results_path:
        raise ValueError("results_path不能为空")
    
    if config.output.log_level not in ['DEBUG', 'INFO', 'WARNING', 'ERROR']:
        raise ValueError(f"不支持的日志级别: {config.output.log_level}")


def config_to_args(config: Config) -> Any:
    """
    将配置对象转换为类似命令行参数的对象
    
    这个函数用于保持与现有代码的兼容性，将新的配置系统转换为
    原有代码期望的参数格式。
    
    Args:
        config (Config): 配置对象
        
    Returns:
        Any: 参数对象
    """
    from types import SimpleNamespace
    
    args = SimpleNamespace()
    
    # 模型配置
    args.model_type = config.model.type
    args.vllm_config = config.model.vllm_config
    args.openai_config = config.model.openai_config
    args.generation_config = config.model.generation_config
    
    # 从generation_config中提取max_tokens
    args.max_tokens = config.model.generation_config.get('max_tokens', 2048)
    
    # 数据配置
    args.dataset_path = config.data.dataset_path
    args.image_root = config.data.image_root
    args.max_samples = config.data.max_samples
    args.field_mapping = config.data.field_mapping
    
    # MCTS配置
    args.num_rollouts = config.mcts.num_rollouts
    args.max_depth_allowed = config.mcts.max_depth_allowed
    args.mcts_exploration_weight = config.mcts.exploration_weight
    args.mcts_weight_scheduler = config.mcts.weight_scheduler
    args.mcts_discount_factor = config.mcts.discount_factor
    args.num_a1_steps = config.mcts.num_a1_steps
    args.num_subquestions = config.mcts.num_subquestions
    args.num_votes = config.mcts.num_votes
    args.mcts_num_last_votes = config.mcts.num_last_votes
    args.disable_a1 = config.mcts.disable_a1
    args.disable_a5 = config.mcts.disable_a5
    args.enable_potential_score = config.mcts.enable_potential_score
    args.modify_prompts_for_rephrasing = config.mcts.modify_prompts_for_rephrasing
    
    # 实验配置
    args.seed = config.experiment.seed
    args.note = config.experiment.note
    args.start_idx = config.experiment.start_idx
    args.end_idx = config.experiment.end_idx
    args.verbose = config.experiment.verbose
    args.save_tree = config.experiment.save_tree
    args.save_intermediate = config.experiment.save_intermediate
    args.max_num_worker = config.experiment.max_num_worker
    args.test_batch_size = config.experiment.test_batch_size
    
    # 输出配置
    args.results_path = config.output.results_path
    args.outputs_dir = config.output.outputs_dir
    args.log_level = config.output.log_level
    args.save_answer_sheets = config.output.save_answer_sheets
    args.save_solutions = config.output.save_solutions
    args.save_statistics = config.output.save_statistics
    
    # 创建与原始GSM8K生成器完全一致的输出目录结构
    import os
    from datetime import datetime
    
    # 构建主输出目录，模拟原始的post_process_args逻辑
    suffix = "---[" + args.note + "]" if args.note else ""
    model_name = "multimodal"  # 多模态模型标识
    dataset_name = "multimodal"  # 数据集标识
    
    # 创建主输出目录：{outputs_dir}/{dataset}/{model}/{timestamp}---[note]
    args.run_outputs_dir = os.path.join(
        args.outputs_dir,
        dataset_name,
        model_name,
        f"{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}" + suffix
    )
    os.makedirs(args.run_outputs_dir, exist_ok=True)
    
    # 创建答案保存子目录：存储每个问题的详细解答结果
    args.answer_sheets_dir = os.path.join(args.run_outputs_dir, "answer_sheets")
    os.makedirs(args.answer_sheets_dir, exist_ok=True)
    
    # 多模态配置
    args.image_format = config.multimodal.image_format
    args.max_image_size = config.multimodal.max_image_size
    args.always_include_image = config.multimodal.always_include_image
    args.image_placeholder = config.multimodal.image_placeholder
    args.show_image_info = config.multimodal.show_image_info
    args.log_image_usage = config.multimodal.log_image_usage
    
    # 设置提示文件路径（使用GSM8K作为默认数据集）
    dataset_name = "GSM8K"
    prompts_dir = f"prompts/{dataset_name}"
    
    args.decompose_template_path = f"{prompts_dir}/decompose/decompose_template.json"
    args.decompose_prompt_path = f"{prompts_dir}/decompose/decompose_prompt.txt"
    args.fewshot_cot_prompt_path = f"{prompts_dir}/fewshot_cot/fewshot_cot_prompt.txt"
    args.fewshot_cot_config_path = f"{prompts_dir}/fewshot_cot/fewshot_cot_config.json"
    args.fewshot_ost_prompt_path = f"{prompts_dir}/fewshot_ost/fewshot_ost_prompt.txt"
    args.fewshot_ost_config_path = f"{prompts_dir}/fewshot_ost/fewshot_ost_config.json"
    
    # 重述相关路径（仅在启用A5时需要）
    if not args.disable_a5:
        args.rephrasing_prompt_template_path = f"{prompts_dir}/rephrasing_prompt_template.txt"
        
        # 根据 modify_prompts_for_rephrasing 参数决定是否使用专门的重述提示
        if args.modify_prompts_for_rephrasing:
            args.decompose_prompt_rephrased_path = f"{prompts_dir}/decompose/decompose_prompt_rephrased.txt"
            args.fewshot_cot_prompt_rephrased_path = f"{prompts_dir}/fewshot_cot/fewshot_cot_prompt_rephrased.txt"
            args.fewshot_ost_prompt_rephrased_path = f"{prompts_dir}/fewshot_ost/fewshot_ost_prompt_rephrased.txt"
        else:
            # 使用标准提示作为重述版本
            args.decompose_prompt_rephrased_path = f"{prompts_dir}/decompose/decompose_prompt.txt"
            args.fewshot_cot_prompt_rephrased_path = f"{prompts_dir}/fewshot_cot/fewshot_cot_prompt.txt"
            args.fewshot_ost_prompt_rephrased_path = f"{prompts_dir}/fewshot_ost/fewshot_ost_prompt.txt"
    
    return args


# 保持向后兼容性
def load_multimodal_config(config_path: str) -> Config:
    """
    加载多模态配置文件（向后兼容函数）
    
    Args:
        config_path (str): 配置文件路径
        
    Returns:
        Config: 配置对象
    """
    return load_config(config_path) 